import sys
sys.path.append("..")
sys.path.append("../..")
from Dataloader.dataloader_utils import Merge_data
from Dataloader.twitterloader import TwitterSet
import torch.nn as nn

from SentModel.Sent2Vec import W2VRDMVec
from PropModel.SeqPropagation import GRUModel
import os
from RumdetecFramework.AdverRumorFramework import EnhancedAdverModel

def obtain_general_set(tr_prefix, dev_prefix, te_prefix):
    tr_set = TwitterSet()
    tr_set.load_data_fast(data_prefix=tr_prefix, min_len=5)
    dev_set = TwitterSet()
    dev_set.load_data_fast(data_prefix=dev_prefix)
    te_set = TwitterSet()
    te_set.load_data_fast(data_prefix=te_prefix)
    return tr_set, dev_set, te_set


log_dir = str(__file__).split(".")[0]
if not os.path.exists(log_dir):
    os.system("mkdir %s"%log_dir)
else:
    os.system("rm -rf %s" % log_dir)
    os.system("mkdir %s" % log_dir)

best_model_file_dic = {
    "ottawashooting": "EnhanceAdverRDM_ottawashooting_0.66.pkl",
    "sydneysiege":"EnhanceAdverRDM_sydneysiege_0.73.pkl",
    "germanwings":"EnhanceAdverRDM_germanwings-crash_0.67.pkl",
    "ferguson":"EnhanceAdverRDM_ferguson_0.76.pkl",
    "charliehebdo":"EnhanceAdverRDM_charliehebdo_0.83.pkl"
}

# for few_samples in [10, 30, 50, 100]:
for t in range(10):
    for i in range(5):
        tr, dev, te = obtain_general_set("../../data/twitter_tr%d"%i, "../../data/twitter_dev%d"%i, "../../data/twitter_te%d"%i)
        tr.filter_short_seq(min_len=5)
        tr.trim_long_seq(10)
        test_event_name = te.data[te.data_ID[0]]['event']
        print("%s : (dev event)/(test event)/(train event) = %3d/%3d/%3d" % (
                    test_event_name, len(dev), len(te), len(tr)
                )
        )
        print("\n\n===========%s Rumor PreTrain===========\n\n"%te.data[te.data_ID[0]]['event'])

        lvec = W2VRDMVec("../../saved/glove_en/", 300, seg=None, emb_update=False)
        prop = GRUModel(300, 256, 1, 0.2)
        cls = nn.Linear(256, 2)
        model = EnhancedAdverModel(lvec, prop, cls,
                         topic_label_num=5, prop_specific_only=True, batch_size=20, grad_accum_cnt=1)
        rdm_log_dir = "%s_RDM"%log_dir
        if not os.path.exists(rdm_log_dir):
            os.system("mkdir %s"%rdm_log_dir)
        # """
        # # Continual Training
        # """
        # # few samples = [10, 30,  50,  100]
        # dev.filter_short_seq(min_len=5)
        # dev.trim_long_seq(10)
        # continual_cnt = 50
        # leakage_frac = continual_cnt*1.0/len(dev)
        # ddev1, _ = dev.split(percent=[leakage_frac, 1.0])
        #
        # leakage_frac = (1000 - continual_cnt)*1.0/len(tr)
        # ttr, _ = tr.split(percent=[leakage_frac, 1.0])
        #
        # dev = te
        # tr = Merge_data(ttr, ddev1)
        #
        # log_dir = "%s_FT" % log_dir
        # if not os.path.exists(log_dir):
        #     os.system("mkdir %s" % log_dir)
        # else:
        #     os.system("rm -rf %s" % log_dir)
        #     os.system("mkdir %s" % log_dir)
        #
        # model.load_model("../../saved/%s"%best_model_file_dic[test_event_name])
        # model.AdverTrain(tr, dev, te,
        #                  Unseen_every=-1, lambda_1=0.9, lambda_2=0.0,
        #                  valid_every=100, max_epochs=3, lr_discount=0.1,
        #                  log_dir= log_dir, log_suffix=test_event_name, model_file="../../saved/EnhanceAdverRDM_FT_%s.pkl"%test_event_name)

        # model.train_iters(tr, dev, te,
        #                     valid_every=100, max_epochs=10, lr_discount=1.0,
        #                     best_valid_acc=0.0, best_test_acc=0.0, best_valid_test_acc=0.0,
        #                     log_dir=rdm_log_dir, log_suffix=test_event_name, model_file="RDM_E.pkl")
        # model.load_model("RDM_E.pkl")
        # os.system("rm RDM_E.pkl")
        model.AdverTrain(tr, dev, te,
                         lambda_1=0.9, lambda_2=0.0, lambda_3=0, # lambda 1: adver, lambda 2:detach, lambda 3: Unseen
                         valid_every=100, max_epochs=10, lr_discount=1.0,
                         log_dir=log_dir, log_suffix=test_event_name, model_file="../../saved/EnhanceAdverRDM_%s.pkl"%test_event_name)
